use bb8::{Pool, PooledConnection};
use bb8_postgres::PostgresConnectionManager;
use chrono::{FixedOffset, NaiveDate, NaiveDateTime, NaiveTime, TimeZone};
use polars::prelude::*;
use std::collections::HashMap;
use thiserror::Error;
use tokio_postgres::NoTls;
use futures::future::join_all;

#[derive(Error, Debug)]
pub enum MyError {
    #[error("数据库错误: {0}")]
    Database(#[from] tokio_postgres::Error),
    #[error("连接池错误: {0}")]
    Pool(String),
    #[error("日期解析错误: {0}")]
    DateParse(#[from] chrono::ParseError),
    #[error("无效表名: {0}")]
    InvalidTableName(String),
    #[error("Polars错误: {0}")]
    Polars(#[from] PolarsError),
    #[error("日期时间解析错误: {0}")]
    DateTimeParseError(String),
}

// 连接池类型别名
type ConnectionPool = Pool<PostgresConnectionManager<NoTls>>;

// 创建连接池
async fn create_connection_pool(conn_str: &str) -> Result<ConnectionPool, MyError> {
    let config = conn_str.parse().map_err(MyError::Database)?;
    let manager = PostgresConnectionManager::new(config, NoTls);
    Pool::builder()
        .max_size(16)
        .min_idle(Some(4))
        .connection_timeout(std::time::Duration::from_secs(30))
        .build(manager)
        .await
        .map_err(|e| MyError::Pool(e.to_string()))
}

// 简化的 DataFrame 构建函数 - 主要修改这里！
fn build_dataframe_from_rows(
    rows: Vec<tokio_postgres::Row>,
    table_name: &str,
) -> Result<DataFrame, MyError> {
    let capacity = rows.len();

    let mut datetime_vec = Vec::with_capacity(capacity);
    let mut timestamp_vec = Vec::with_capacity(capacity);
    let mut open_vec = Vec::with_capacity(capacity);
    let mut high_vec = Vec::with_capacity(capacity);
    let mut low_vec = Vec::with_capacity(capacity);
    let mut close_vec = Vec::with_capacity(capacity);
    let mut volume_vec = Vec::with_capacity(capacity);

    // 北京时间偏移
    let beijing_offset = FixedOffset::east_opt(8 * 3600)
        .ok_or_else(|| MyError::DateTimeParseError("Invalid timezone offset".to_string()))?;

    for row in rows {
        let datetime_utc: NaiveDateTime  = row.get(0);

        // 调整时间戳(如果需要)
        let adjusted_datetime = adjust_timestamp_by_period(datetime_utc, table_name);

        // 转换为北京时间并获取时间戳
        let datetime_beijing = beijing_offset
            .from_local_datetime(&adjusted_datetime)
            .single()
            .ok_or_else(|| MyError::DateTimeParseError("Ambiguous datetime".to_string()))?;

        timestamp_vec.push(datetime_beijing.timestamp_millis());
        datetime_vec.push(adjusted_datetime);
        open_vec.push(row.get::<_, f64>(1));
        high_vec.push(row.get::<_, f64>(2));
        low_vec.push(row.get::<_, f64>(3));
        close_vec.push(row.get::<_, f64>(4));
        volume_vec.push(row.get::<_, f64>(5));
    }

    let df = DataFrame::new(vec![
        Column::new("datetime".into(), datetime_vec),
        Column::new("timestamps".into(), timestamp_vec),
        Column::new("open".into(), open_vec),
        Column::new("high".into(), high_vec),
        Column::new("low".into(), low_vec),
        Column::new("close".into(), close_vec),
        Column::new("volume".into(), volume_vec),
    ])?;

    Ok(df)
}

fn adjust_timestamp_by_period(dt: NaiveDateTime, table_name: &str) -> NaiveDateTime {
    if let Some(period_str) = table_name.split('_').last() {
        if period_str == "86400" && dt.time() == NaiveTime::from_hms_opt(0, 0, 0).unwrap() {
            return dt.date().and_hms_opt(15, 0, 0).unwrap();
        }
    }
    dt
}

// 处理连接获取错误的辅助函数
async fn get_connection_from_pool(
    pool: &ConnectionPool,
) -> Result<PooledConnection<'_, PostgresConnectionManager<NoTls>>, MyError> {
    pool.get().await.map_err(|e| MyError::Pool(e.to_string()))
}
fn calculate_batch_size(table_count: usize) -> usize {
    // 1. 基于数据库连接池大小(最重要)
    let max_connections = 16; // 您的连接池大小,这是一个经验值
    let recommended = (max_connections * 2 / 3).max(4).min(12);

    // 2. 基于表数量
    let based_on_tables = match table_count {
        0..=10 => 4,     // 表少时并发少
        11..=30 => 6,    // 中等
        31..=50 => 8,    // 较多
        51..=100 => 10,  // 大量表
        101..=150 => 12, // 超多表
        _ => 12,         // 极限情况
    };

    // 取较小值,避免超过连接池限制
    recommended.min(based_on_tables)
}
pub async fn load_data_parallel_optimized(
    table_names: Vec<String>,
    s_date_str: &str,
    e_date_str: &str,
) -> Result<HashMap<String, DataFrame>, MyError> {
    let conn_str = "postgresql://postgres:123456@localhost/Future";

    // 验证表名
    for table_name in &table_names {
        if !table_name.chars().all(|c| c.is_alphanumeric() || c == '_') {
            return Err(MyError::InvalidTableName(table_name.to_string()));
        }
    }

    // 解析日期 - 现在直接使用 TIMESTAMP 参数
    let s_date = NaiveDate::parse_from_str(s_date_str, "%Y%m%d")?;
    let e_date = NaiveDate::parse_from_str(e_date_str, "%Y%m%d")?;

    // 创建 TIMESTAMP 参数
    let s_date_bound = s_date.and_hms_opt(0, 0, 0)
        .ok_or_else(|| MyError::DateTimeParseError("Invalid start date".to_string()))?;
    let e_date_bound = e_date.and_hms_opt(23, 59, 59)
        .ok_or_else(|| MyError::DateTimeParseError("Invalid end date".to_string()))?;

    // 使用连接池
    let pool = create_connection_pool(conn_str).await?;
    let mut result_map = HashMap::new();

    let batch_size = calculate_batch_size(table_names.len());
    
    for (_batch_num, chunk) in table_names.chunks(batch_size).enumerate() {
        let mut batch_tasks = Vec::new();

        for table_name in chunk {
            let table_name = table_name.clone();
            let pool = pool.clone();
            let s_date_bound = s_date_bound;
            let e_date_bound = e_date_bound;

            batch_tasks.push(async move {
                let conn = get_connection_from_pool(&pool).await?;
                let query = format!(
                    "SELECT datetime, open::double precision, high::double precision,
                            low::double precision, close::double precision, volume::double precision
                     FROM {}
                     WHERE datetime BETWEEN $1 AND $2
                     ORDER BY datetime ASC",
                    table_name
                );

                // 直接传递 NaiveDateTime 参数
                let rows = conn.query(&query, &[&s_date_bound, &e_date_bound]).await?;
                let df = build_dataframe_from_rows(rows, &table_name)?;

                Ok::<(String, DataFrame), MyError>((table_name, df))
            });
        }

        let batch_results = join_all(batch_tasks).await;
        for result in batch_results {
            match result {
                Ok((table_name, df)) => {
                    result_map.insert(table_name, df);
                }
                Err(e) => return Err(e),
            }
        }

        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
    }

    Ok(result_map)
}